#!/usr/bin/env python
# -*- coding: utf-8 -*-

import os
import sys
import json
import subprocess
import argparse
import datetime
import shutil
import logging
import re
from collections import defaultdict
import copy
import glob

# --- 全局配置 ---
# 注意：此脚本需要与 lib/ 和 mysql_configs/ 目录在同一级别
LIB_DIR = "lib"
# 使用与 sysbench 相同的经过验证的 MySQL 设置脚本
SETUP_SCRIPT = os.path.join(LIB_DIR, "setup_mysql_instances.sh")
# [NEW] 指向新的 BenchmarkSQL 运行脚本
BENCHMARKSQL_SCRIPT = os.path.join(LIB_DIR, "run_benchmarksql.sh")
TEMP_CONFIG_DIR = "temp_configs"
TEMP_PROPS_DIR = "temp_benchmarksql_props" # [NEW] 存放临时的 props 文件

logging.basicConfig(level=logging.INFO,
                    format='%(asctime)s - %(levelname)s - %(message)s',
                    stream=sys.stdout)

# --- 辅助函数 (与 sysbench 版本基本相同，但做微调) ---

def run_command(command, log_file):
    """执行命令并记录输出。兼容 Python 2/3。"""
    logging.info("执行命令: {}".format(' '.join(command)))
    try:
        process = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, universal_newlines=True, bufsize=1)
        with open(log_file, 'a') as f:
            for line in process.stdout:
                sys.stdout.write(line)
                f.write(line)
        process.wait()
        if process.returncode != 0:
            logging.error("命令执行失败，返回码: {}".format(process.returncode))
            return False
        logging.info("命令执行成功.")
        return True
    except Exception as e:
        logging.error("执行命令时发生异常: {}".format(e))
        return False

def generate_mysql_config(base_configs, mysql_params, output_path):
    """生成 MySQL 配置文件。"""
    try:
        os.makedirs(os.path.dirname(output_path), exist_ok=True)
        with open(output_path, 'w') as f:
            f.write("# Auto-generated MySQL config for this test run\n\n")
            for base_path in base_configs:
                f.write("!include {}\n".format(os.path.abspath(base_path)))
            f.write("\n[mysqld]\n")
            for key, value in mysql_params.items():
                f.write("{:<35} = {}\n".format(key, value))
        logging.info("已生成临时MySQL配置文件: {}".format(output_path))
        return True
    except Exception as e:
        logging.error("生成MySQL配置文件失败: {}".format(e))
        return False

# [NEW] 生成 BenchmarkSQL 属性文件
def generate_benchmarksql_props(base_props, benchmarksql_params, output_path):
    """根据基础配置和待测参数生成 BenchmarkSQL 的属性文件。"""
    try:
        os.makedirs(os.path.dirname(output_path), exist_ok=True)
        # 深拷贝以避免修改原始配置字典
        full_props = copy.deepcopy(base_props)
        # 将待优化的参数更新到 props 字典中
        full_props.update(benchmarksql_params)

        with open(output_path, 'w') as f:
            f.write("# Auto-generated BenchmarkSQL properties for this test run\n")
            for key, value in sorted(full_props.items()):
                 f.write("{}={}\n".format(key, value))
        logging.info("已生成 BenchmarkSQL 属性文件: {}".format(output_path))
        return True
    except Exception as e:
        logging.error("生成 BenchmarkSQL 属性文件失败: {}".format(e))
        return False

def apply_os_params(os_params, devices):
    """应用操作系统参数。"""
    original_values = {}
    logging.info("正在应用操作系统参数...")
    for key, value in os_params.items():
        try:
            if key.startswith("vm.") or key.startswith("net."):
                original = subprocess.check_output(['sysctl', '-n', key]).decode('utf-8').strip()
                original_values[key] = original
                logging.info("  - 设置 {} = {} (原值: {})".format(key, value, original))
                subprocess.run(['sudo', 'sysctl', '-w', "{}={}".format(key, value)], check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
            elif key.startswith('/sys/'):
                for device_path in devices:
                    dev_name = os.path.basename(device_path)
                    actual_path = key.replace('nvme0n1', dev_name)
                    if os.path.exists(actual_path):
                        with open(actual_path, 'r') as f:
                            raw_original = f.read().strip()
                            match = re.search(r'\[(\w+-?\w*)\]', raw_original)
                            original = match.group(1) if match else raw_original.split()[0]
                        if key not in original_values: original_values[key] = original
                        logging.info("  - 设置 {} = {} (原值: {})".format(actual_path, value, original))
                        with open(actual_path, 'w') as f: f.write(str(value))
        except Exception as e: logging.error("设置OS参数 {} 失败: {}".format(key, e))
    return original_values

def revert_os_params(original_values, devices):
    """恢复原始操作系统参数。"""
    if not original_values: return
    logging.info("正在恢复原始操作系统参数...")
    for key, value in original_values.items():
        try:
            if key.startswith("vm.") or key.startswith("net."):
                subprocess.run(['sudo', 'sysctl', '-w', "{}={}".format(key, value)], check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
            elif key.startswith('/sys/'):
                for device_path in devices:
                    dev_name = os.path.basename(device_path)
                    actual_path = key.replace('nvme0n1', dev_name)
                    if os.path.exists(actual_path):
                        with open(actual_path, 'w') as f: f.write(str(value))
        except Exception as e: logging.error("恢复OS参数 {} 失败: {}".format(key, e))

# [MODIFIED] 解析 BenchmarkSQL 的结果
def parse_benchmarksql_results(base_output_dir, num_instances):
    """解析所有实例的 BenchmarkSQL TpmC 结果并求和。"""
    total_tpmc = 0.0
    start_port = 3307
    
    logging.info("正在解析 BenchmarkSQL 结果...")
    for i in range(num_instances):
        port = start_port + i
        instance_result_dir = os.path.join(base_output_dir, "port_{}".format(port))
        
        # 寻找最新的 my_result_* 目录
        list_of_dirs = glob.glob(os.path.join(instance_result_dir, 'my_result_*'))
        if not list_of_dirs:
            logging.warning("在 {} 中未找到 BenchmarkSQL 结果目录。".format(instance_result_dir))
            continue
        latest_dir = max(list_of_dirs, key=os.path.getmtime)
        summary_file = os.path.join(latest_dir, 'summary.txt')

        if not os.path.exists(summary_file):
            logging.warning("在 {} 中未找到 summary.txt。".format(latest_dir))
            continue
        
        try:
            with open(summary_file, 'r') as f:
                content = f.read()
                # 匹配 "(New-Order) TpmC:" 行
                match = re.search(r"\(New-Order\)\s+TpmC:\s*([\d\.]+)", content)
                if match:
                    tpmc = float(match.group(1))
                    total_tpmc += tpmc
                    logging.info("  - 实例 (端口 {}) TpmC: {:.2f}".format(port, tpmc))
                else:
                    logging.warning("在 {} 中未能解析到 TpmC 值。".format(summary_file))
        except Exception as e:
            logging.error("解析文件 {} 失败: {}".format(summary_file, e))
            
    return total_tpmc

# [MODIFIED] 生成最终报告
def generate_final_report(final_best_results, output_dir):
    """生成最终的优化报告。"""
    report_path = os.path.join(output_dir, "final_analysis_report.txt")
    logging.info("正在生成最终报告: {}".format(report_path))
    with open(report_path, 'w') as f:
        f.write("="*80 + "\n")
        f.write("=      BenchmarkSQL Auto-Tuning Final Report (Control Variable Method)     =\n")
        f.write("="*80 + "\n\n")

        # BenchmarkSQL 只有一个核心 workload
        workload = "tpcc"
        if workload in final_best_results:
            best_run = final_best_results[workload]
            f.write("--- Best Configuration for TPC-C Workload ---\n")
            f.write("  [+] Achieved Best Total TpmC (New-Order): {:.2f}\n".format(best_run['total_tpmc']))
            f.write("  [+] Found after {} test runs in {} epochs.\n\n".format(best_run['total_runs'], best_run['epoch']))
            f.write("  Optimal Parameters Found:\n")
            f.write("  " + "-"*40 + "\n")
            params = best_run['params']
            
            f.write("    MySQL Parameters:\n")
            if params.get('mysql'):
                for k, v in sorted(params['mysql'].items()):
                    f.write("        {:<30} = {}\n".format(k, v))
            else:
                f.write("        (No MySQL parameters tuned)\n")

            f.write("\n    OS Parameters:\n")
            if params.get('os'):
                for k, v in sorted(params['os'].items()):
                    f.write("        {:<30} = {}\n".format(k, v))
            else:
                f.write("        (No OS parameters tuned)\n")

            f.write("\n    BenchmarkSQL Parameters:\n")
            if params.get('benchmarksql'):
                for k, v in sorted(params['benchmarksql'].items()):
                    f.write("        {:<30} = {}\n".format(k, v))
            else:
                f.write("        (No BenchmarkSQL parameters tuned)\n")
            
            f.write("\n" + "="*80 + "\n\n")

    with open(report_path, 'r') as f:
        print("\n" + f.read())

def run_single_test(test_id, params_to_test, config, run_log_file, needs_load):
    """
    封装单次 BenchmarkSQL 测试逻辑。
    - needs_load: bool, 指示是否需要运行 runDatabaseBuild.sh
    """
    mysql_inst_cfg = config['mysql_instances']
    devices = mysql_inst_cfg['devices'].split(',')
    num_instances = mysql_inst_cfg['count']
    ports_str = ",".join(map(str, range(3307, 3307 + num_instances)))

    logging.info("--- [Test ID: {}] Running BenchmarkSQL test (Load Data: {}) ---".format(test_id, needs_load))
    logging.info("Parameters: {}".format(json.dumps(params_to_test)))

    original_os_params = {}
    try:
        # 1. 应用操作系统参数
        original_os_params = apply_os_params(params_to_test.get('os', {}), devices)

        # 2. 生成 MySQL 配置文件并重启实例
        temp_mysql_config_path = os.path.join(TEMP_CONFIG_DIR, "config_{}.cnf".format(test_id))
        if not generate_mysql_config(config['mysql_base_configs'], params_to_test['mysql'], temp_mysql_config_path):
             raise Exception("Failed to generate MySQL config")

        setup_cmd = ['sudo', SETUP_SCRIPT, '-n', str(num_instances), '-c', mysql_inst_cfg['numa_nodes'], '-D', mysql_inst_cfg['devices'], '-f', os.path.abspath(temp_mysql_config_path)]
        if not run_command(setup_cmd, run_log_file):
            raise Exception("MySQL setup script failed")

        # 3. 生成 BenchmarkSQL 属性文件
        temp_props_path = os.path.join(TEMP_PROPS_DIR, "props_{}.properties".format(test_id))
        if not generate_benchmarksql_props(config['benchmarksql_base_props'], params_to_test['benchmarksql'], temp_props_path):
            raise Exception("Failed to generate BenchmarkSQL properties")

        # 4. 构建并执行 BenchmarkSQL 命令
        # run_benchmarksql.sh 脚本将负责处理并行和数据加载逻辑
        run_cmd = ['sudo', BENCHMARKSQL_SCRIPT, '-p', ports_str, '-f', os.path.abspath(temp_props_path)]
        if needs_load:
            run_cmd.append('--load')
        
        if not run_command(run_cmd, run_log_file):
            raise Exception("BenchmarkSQL script failed")

        # 5. 解析结果
        total_tpmc = parse_benchmarksql_results('benchmarksql_results', num_instances)
        logging.info("RESULT for Test ID {}: Total TpmC: {:.2f}".format(test_id, total_tpmc))

    except Exception as e:
        logging.error("Test run {} failed: {}".format(test_id, e))
        total_tpmc = 0.0
    finally:
        revert_os_params(original_os_params, devices)

    return total_tpmc


def main():
    parser = argparse.ArgumentParser(
        description="BenchmarkSQL 自动化性能优化器 (v1.0)",
        formatter_class=argparse.RawTextHelpFormatter
    )
    parser.add_argument('config_file', help="JSON格式的 BenchmarkSQL 优化配置文件路径。")
    args = parser.parse_args()

    # --- 1. 加载和解析配置 ---
    if not os.path.exists(args.config_file):
        logging.error("配置文件未找到: {}".format(args.config_file))
        sys.exit(1)
    with open(args.config_file, 'r') as f:
        config = json.load(f)

    # 获取触发数据重建的参数列表
    rebuild_triggers = set(config.get('rebuild_trigger_params', []))

    # --- 2. 设置输出目录 ---
    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    top_level_dir = "benchmarksql_run_{}".format(timestamp)
    os.makedirs(top_level_dir, exist_ok=True)
    os.makedirs(TEMP_CONFIG_DIR, exist_ok=True)
    os.makedirs(TEMP_PROPS_DIR, exist_ok=True)
    shutil.copy(args.config_file, os.path.join(top_level_dir, "run_config.json"))
    run_log_file = os.path.join(top_level_dir, "full_run.log")
    file_handler = logging.FileHandler(run_log_file)
    file_handler.setFormatter(logging.Formatter('%(asctime)s - %(levelname)s - %(message)s'))
    logging.getLogger().addHandler(file_handler)
    logging.info("所有输出将保存在: {}".format(top_level_dir))

    # --- 3. 构建参数列表 ---
    params_to_optimize = []
    for group_name, params in config['optimization_parameters'].items():
        for param in params:
            if len(param['values']) > 1:
                params_to_optimize.append({'group': group_name, 'name': param['name'], 'values': param['values']})
    logging.info("优化策略: 控制变量法 (智能数据加载)")
    logging.info("数据加载触发参数: {}".format(rebuild_triggers or 'None'))

    # --- 4. 核心优化循环 ---
    final_best_results = {}
    workload = "tpcc"  # BenchmarkSQL 只有一个核心 workload

    logging.info("============== 开始为 Workload: [{}] 进行优化 ==============".format(workload))
    
    # 确定初始基线配置
    best_params_for_workload = defaultdict(dict)
    for group_name, params in config['optimization_parameters'].items():
        for param in params:
            best_params_for_workload[group_name][param['name']] = param['values'][0]
    
    best_tpmc_for_workload = 0.0
    total_runs_for_workload = 0

    for epoch in range(1, config.get('optimization_epochs', 1) + 1):
        logging.info("--- [Workload: {}, Epoch: {}] ---".format(workload, epoch))
        
        # Epoch 开始时的基线测试，总是需要 load data
        if epoch == 1:
            logging.info("--- Running initial baseline test ---")
            test_id = "{}-E{}-baseline".format(workload, epoch)
            # 第一次运行，needs_load=True
            baseline_tpmc = run_single_test(test_id, best_params_for_workload, config, run_log_file, needs_load=True)
            best_tpmc_for_workload = baseline_tpmc
            total_runs_for_workload += 1
            logging.info("Baseline TpmC for Epoch {}: {:.2f}".format(epoch, best_tpmc_for_workload))

        # OVAT 核心循环
        for param_info in params_to_optimize:
            group = param_info['group']
            name = param_info['name']
            values = param_info['values']
            
            param_fullname = "{}.{}".format(group, name)
            current_param_needs_rebuild = param_fullname in rebuild_triggers
            
            logging.info("--- Optimizing parameter: [{}] (Reload required: {}) ---".format(param_fullname, current_param_needs_rebuild))
            
            current_best_value = best_params_for_workload[group][name]
            
            for value in values:
                if str(value) == str(current_best_value):
                    continue

                params_to_test = copy.deepcopy(best_params_for_workload)
                params_to_test[group][name] = value
                
                safe_name = name.replace('/', '_')
                test_id = "{}-E{}-{}.{}-{}".format(workload, epoch, group, safe_name, value)
                
                # 传递 needs_load 标志
                current_tpmc = run_single_test(test_id, params_to_test, config, run_log_file, needs_load=current_param_needs_rebuild)
                total_runs_for_workload += 1

                if current_tpmc > best_tpmc_for_workload:
                    logging.info("*** NEW BEST FOUND for [{}]! TpmC: {:.2f} -> {:.2f} ***".format(workload, best_tpmc_for_workload, current_tpmc))
                    logging.info("*** Parameter '{}' updated to '{}' ***".format(param_fullname, value))
                    best_tpmc_for_workload = current_tpmc
                    best_params_for_workload = copy.deepcopy(params_to_test)
                    current_best_value = value
        
        logging.info("--- Epoch {} for workload [{}] finished. Best TpmC so far: {:.2f} ---".format(epoch, workload, best_tpmc_for_workload))

    # 记录最终结果
    final_best_results[workload] = {
        'total_tpmc': best_tpmc_for_workload,
        'params': best_params_for_workload,
        'epoch': config.get('optimization_epochs', 1),
        'total_runs': total_runs_for_workload
    }

    # --- 5. 清理和报告 ---
    if os.path.exists("benchmarksql_results"):
        shutil.move("benchmarksql_results", os.path.join(top_level_dir, "benchmarksql_results"))
    if os.path.exists(TEMP_CONFIG_DIR):
        shutil.move(TEMP_CONFIG_DIR, os.path.join(top_level_dir, "generated_mysql_configs"))
    if os.path.exists(TEMP_PROPS_DIR):
        shutil.move(TEMP_PROPS_DIR, os.path.join(top_level_dir, "generated_benchmarksql_props"))
    
    generate_final_report(final_best_results, top_level_dir)
    logging.info("优化过程完成。所有数据已保存在: {}".format(top_level_dir))

if __name__ == '__main__':
    if os.geteuid() != 0:
        print("错误: 此脚本需要root权限来修改系统参数和管理服务。请使用 'sudo python ...' 运行。")
        sys.exit(1)
    main()